/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: op remote update execute.
 * Author: qingjicheng
 * Create: 2022-12-15
 */

#ifndef OP_COLL_REMOTE_UPDATE_EXECUTOR_H
#define OP_COLL_REMOTE_UPDATE_EXECUTOR_H

#include <map>
#include <unordered_map>
#include <atomic>
#include <memory>
#include <thread>
#include <vector>
#include <condition_variable>
#include "op_executor.h"
#include "transport_heterog_pub.h"
#include "hccl_operator.h"
#include "network_manager_pub.h"
#include "transport_heterog_pub.h"
#include "unsafe_table.h"

namespace hccl {
constexpr u32 REDUCE_THREAD_NUM = 8;

class OpCollRemoteUpdateExecutor : public OpExecutor {
    using USTable = Hss::UnsafeTable<s64, u32>;
    using InsertResult = typename USTable::InsertResult;
public:
    explicit OpCollRemoteUpdateExecutor();
    ~OpCollRemoteUpdateExecutor() override;
    HcclResult Init(std::map<u32, std::unique_ptr<TransportHeterog>> &transportHandleMap,
        std::map<u32, u32> &psRankIdMap, std::vector<u32> &workerList, u32 threadNum = REDUCE_THREAD_NUM,
        bool needThread = false);
    HcclResult ExecuteOp(HcclOpPtr opPtr) override;
    void PsUpdatePrepare(u32 keySize);
    void ClearAllTask(HcclOperator<EmbeddingServiceParam> *opPtr);

private:
    HcclResult TaskSend(OpTaskPtr &task);
    HcclResult TaskRecv(OpTaskPtr &task);
    HcclResult TaskIsend(OpTaskPtr &task);
    HcclResult TaskImrecv(OpTaskPtr &task);
    HcclResult TaskWait(OpTaskPtr &task);
    HcclResult TaskProbe(OpTaskPtr &task);
    HcclResult TaskUpdate(OpTaskPtr &task);
    HcclResult TaskMemcpy(OpTaskPtr &task);

private:
    HcclResult DeInit();

    HcclResult ExecuteTask(HcclOperator<EmbeddingServiceParam> *opPtr);

    HcclResult WaitAll();
    HcclResult WaitSome(WaitSomeTask *taskPtr);
    HcclResult WaitRequestList(HcclRequestInfo *requestList, s32 &flag);

    HcclResult ReduceSum(PsReduceSumTask *params);
    void VectorReduce(volatile PsReduceSumTask *params, u32 serialNum);

    void ReduceThread(u32 serialNum);

    std::map<OpTaskType, std::function<HcclResult(OpTaskPtr &)>> taskExeFuncTable_;
    NetworkManager *netResource_ = nullptr;
    std::map<u32, TransportHeterog *> rankTransportMap_;
    std::map<TransportHeterog *, std::queue<HcclMessageInfo *>> transportMsgTempMap_;
    std::vector<HcclRequestInfo *> tempRequestQue_;
    std::map<u32, u32> psRankIdMap_;
    std::atomic<bool> blockFlag_ = {true};
    std::vector<std::unique_ptr<std::thread>> reduceSumThreads_;
    std::condition_variable reduceCv_;
    std::condition_variable reduceEndCv_;
    volatile std::atomic<bool> reduceShutDown_ = {false};
    volatile bool *reduceRunningflags_ = nullptr;
    std::vector<u32> reduceThreadIds_;
    u32 reduceThreadNum_ = 0;

    std::vector<USTable*> realKeyMapForShard_;
    volatile PsReduceSumTask *tempReduceTaskPtr_ = nullptr;
    u32 compThreadNum_ = 0;
    HostMem tempKeys_;
    HostMem tempValues_;
    std::mutex cvMutex_;

    std::mutex reduceMutex_;

    // index workerId index threadId index 无意义 value keyIndex
    std::vector<std::vector<std::vector<u32>>> threadkeyNumForShard_;
    std::vector<u32> keyToThreadId_; // 大小为1024的随机数数组
    std::atomic<u32> realKeyPos_ = {0};
    u32 rankSize_ = 0;
};

} // namespace hccl
#endif // OP_COLL_REMOTE_UPDATE_EXECUTOR_H
